{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Count Number of Possible Root Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #depth-first-search #hash-table #dynamic-programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #深度优先搜索 #哈希表 #动态规划"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: rootCount"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #统计可能的树根数目"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>Alice 有一棵 <code>n</code> 个节点的树，节点编号为 <code>0</code> 到 <code>n - 1</code> 。树用一个长度为 <code>n - 1</code> 的二维整数数组 <code>edges</code> 表示，其中 <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code> ，表示树中节点 <code>a<sub>i</sub></code> 和 <code>b<sub>i</sub></code> 之间有一条边。</p>\n",
    "\n",
    "<p>Alice 想要 Bob 找到这棵树的根。她允许 Bob 对这棵树进行若干次 <strong>猜测</strong> 。每一次猜测，Bob 做如下事情：</p>\n",
    "\n",
    "<ul>\n",
    "\t<li>选择两个 <strong>不相等</strong>&nbsp;的整数&nbsp;<code>u</code> 和&nbsp;<code>v</code>&nbsp;，且树中必须存在边&nbsp;<code>[u, v]</code>&nbsp;。</li>\n",
    "\t<li>Bob 猜测树中&nbsp;<code>u</code>&nbsp;是&nbsp;<code>v</code>&nbsp;的 <strong>父节点</strong>&nbsp;。</li>\n",
    "</ul>\n",
    "\n",
    "<p>Bob 的猜测用二维整数数组&nbsp;<code>guesses</code> 表示，其中&nbsp;<code>guesses[j] = [u<sub>j</sub>, v<sub>j</sub>]</code>&nbsp;表示 Bob 猜&nbsp;<code>u<sub>j</sub></code> 是&nbsp;<code>v<sub>j</sub></code>&nbsp;的父节点。</p>\n",
    "\n",
    "<p>Alice 非常懒，她不想逐个回答&nbsp;Bob 的猜测，只告诉 Bob 这些猜测里面 <strong>至少</strong>&nbsp;有&nbsp;<code>k</code>&nbsp;个猜测的结果为&nbsp;<code>true</code>&nbsp;。</p>\n",
    "\n",
    "<p>给你二维整数数组 <code>edges</code>&nbsp;，Bob 的所有猜测和整数&nbsp;<code>k</code>&nbsp;，请你返回可能成为树根的&nbsp;<strong>节点数目</strong>&nbsp;。如果没有这样的树，则返回 <code>0</code>。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2022/12/19/ex-1.png\" style=\"width: 727px; height: 250px;\" /></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3\n",
    "<b>输出：</b>3\n",
    "<b>解释：</b>\n",
    "根为节点 0 ，正确的猜测为 [1,3], [0,1], [2,4]\n",
    "根为节点 1 ，正确的猜测为 [1,3], [1,0], [2,4]\n",
    "根为节点 2 ，正确的猜测为 [1,3], [1,0], [2,4]\n",
    "根为节点 3 ，正确的猜测为 [1,0], [2,4]\n",
    "根为节点 4 ，正确的猜测为 [1,3], [1,0]\n",
    "节点 0 ，1 或 2 为根时，可以得到 3 个正确的猜测。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2022/12/19/ex-2.png\" style=\"width: 600px; height: 303px;\" /></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1\n",
    "<b>输出：</b>5\n",
    "<b>解释：</b>\n",
    "根为节点 0 ，正确的猜测为 [3,4]\n",
    "根为节点 1 ，正确的猜测为 [1,0], [3,4]\n",
    "根为节点 2 ，正确的猜测为 [1,0], [2,1], [3,4]\n",
    "根为节点 3 ，正确的猜测为 [1,0], [2,1], [3,2], [3,4]\n",
    "根为节点 4 ，正确的猜测为 [1,0], [2,1], [3,2]\n",
    "任何节点为根，都至少有 1 个正确的猜测。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>edges.length == n - 1</code></li>\n",
    "\t<li><code>2 &lt;= n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>1 &lt;= guesses.length &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub>, u<sub>j</sub>, v<sub>j</sub> &lt;= n - 1</code></li>\n",
    "\t<li><code>a<sub>i</sub> != b<sub>i</sub></code></li>\n",
    "\t<li><code>u<sub>j</sub> != v<sub>j</sub></code></li>\n",
    "\t<li><code>edges</code>&nbsp;表示一棵有效的树。</li>\n",
    "\t<li><code>guesses[j]</code>&nbsp;是树中的一条边。</li>\n",
    "\t<li><code>guesses</code>&nbsp;是唯一的。</li>\n",
    "\t<li><code>0 &lt;= k &lt;= guesses.length</code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [count-number-of-possible-root-nodes](https://leetcode.cn/problems/count-number-of-possible-root-nodes/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [count-number-of-possible-root-nodes](https://leetcode.cn/problems/count-number-of-possible-root-nodes/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[[0,1],[1,2],[1,3],[4,2]]\\n[[1,3],[0,1],[1,0],[2,4]]\\n3', '[[0,1],[1,2],[2,3],[3,4]]\\n[[1,0],[3,4],[2,1],[3,2]]\\n1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges)+1\n",
    "        guess_map = {(x, y) for x, y in guesses}\n",
    "\n",
    "        edges_map = [[] for i in range(n)]\n",
    "        for e in edges:\n",
    "            edges_map[e[0]].append(e[1])\n",
    "            edges_map[e[1]].append(e[0])\n",
    "            # if edges_map.get(e[0]):\n",
    "            #     edges_map[e[0]].append(e[1])\n",
    "            # else:\n",
    "            #     edges_map[e[0]] = [e[1]]\n",
    "            \n",
    "            # if edges_map.get(e[1]):\n",
    "            #     edges_map[e[1]].append(e[0])\n",
    "            # else:\n",
    "            #     edges_map[e[1]] = [e[0]]\n",
    "\n",
    "        dp = [0 for i in range(n)]\n",
    "\n",
    "        # 以0为根 算guess数量\n",
    "        # father = [-1 for i in range(n)]\n",
    "        bfs = [0]\n",
    "        visited = [False for i in range(n)]\n",
    "        visited[0] = True\n",
    "        while len(bfs) != 0:\n",
    "            # print(bfs)\n",
    "            layer = []\n",
    "            for x in bfs:\n",
    "                for y in edges_map[x]:\n",
    "                    if visited[y]:\n",
    "                        continue\n",
    "                    else:\n",
    "                        layer.append(y)\n",
    "                        visited[y] = True\n",
    "                        # father[y] = x\n",
    "                        if (x, y) in guess_map:\n",
    "                            dp[0] += 1\n",
    "            bfs = layer\n",
    "        \n",
    "        # dp = [0 for i in range(n)]\n",
    "        # for guess in guesses:\n",
    "        #     if father[guess[1]] == guess[0]:\n",
    "        #         dp[0] += 1\n",
    "\n",
    "        count = 0\n",
    "        bfs = [0]\n",
    "        # visited = [False for i in range(n)]\n",
    "        visited[0] = False\n",
    "        while len(bfs) != 0:\n",
    "            layer = []\n",
    "            for x in bfs:\n",
    "                if dp[x] >= k:\n",
    "                    count += 1\n",
    "                for y in edges_map[x]:\n",
    "                    if not visited[y]:\n",
    "                        continue\n",
    "                    else:\n",
    "                        layer.append(y)\n",
    "                        visited[y] = False\n",
    "                        dp[y] = dp[x]\n",
    "                        if (x, y) in guess_map:\n",
    "                            dp[y] -= 1\n",
    "                        if (y, x) in guess_map:\n",
    "                            dp[y] += 1\n",
    "                        # if guess_map[x][y] == guess_map[y][x]:\n",
    "                        #     dp[y] = dp[x]\n",
    "                        # elif guess_map[x][y]:\n",
    "                        #     dp[y] = dp[x] - 1\n",
    "                        # elif guess_map[y][x]:\n",
    "                        #     dp[y] = dp[x] + 1         \n",
    "            bfs = layer\n",
    "        \n",
    "        # count = 0\n",
    "        # for i in range(n):\n",
    "        #     print(dp[i])\n",
    "        #     if dp[i] >= k :\n",
    "        #         count += 1\n",
    "        return count\n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "                \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n, C = len(edges) + 1, set(map(tuple, guesses))\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        parent, dp, cnt = [-1] * n, [0] * n, 0\n",
    "        order = [0]\n",
    "        stack = [0]\n",
    "        while stack:\n",
    "            u = stack.pop()\n",
    "            for v in g[u]:\n",
    "                if v == parent[u]:\n",
    "                    continue\n",
    "                parent[v] = u\n",
    "                order.append(v)\n",
    "                stack.append(v)\n",
    "        \n",
    "        for u in order[::-1]:\n",
    "            for v in g[u]:\n",
    "                if parent[u] == v:\n",
    "                    continue\n",
    "                if (u, v) in C: dp[u] += 1\n",
    "                dp[u] += dp[v]\n",
    "\n",
    "        for u in order:\n",
    "            if dp[u] >= k: cnt += 1\n",
    "            for v in g[u]:\n",
    "                if parent[u] == v:\n",
    "                    continue\n",
    "                dp[v] = dp[u]\n",
    "                if (u, v) in C: dp[v] -= 1\n",
    "                if (v, u) in C: dp[v] += 1\n",
    "\n",
    "        return cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        nxt = defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            nxt[u].append(v)\n",
    "            nxt[v].append(u)\n",
    "        g = {(u,v) for u,v in guesses}\n",
    "        stack, w = [(0,-1)], 0\n",
    "        while stack:\n",
    "            u,f = stack.pop()\n",
    "            for v in nxt[u]:\n",
    "                if v!=f:\n",
    "                    w += (u,v) in g\n",
    "                    stack.append((v,u)) \n",
    "        res = 0\n",
    "        stack = [(0, -1, w)]\n",
    "        while stack:\n",
    "            u,f,w = stack.pop()\n",
    "            res += w>=k\n",
    "            for v in nxt[u]:\n",
    "                if v!=f:\n",
    "                    w2 = w-((u,v)in g)+((v,u) in g)\n",
    "                    stack.append((v,u,w2))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        dc = defaultdict(list)\n",
    "        for i,j in edges:\n",
    "            dc[i].append(j)\n",
    "            dc[j].append(i)\n",
    "        n = len(edges) + 1\n",
    "        has = [None] * n\n",
    "        dp = [0]\n",
    "        has[0] = 0\n",
    "        c = Counter((i,j) for i,j in guesses)\n",
    "        res = 0\n",
    "        while dp:\n",
    "            x = dp.pop()\n",
    "            for y in dc[x]:\n",
    "                if has[y] == None:\n",
    "                    has[y] = has[x] - c[(x,y)] + c[(y,x)]\n",
    "                    res += c[(x,y)]\n",
    "                    dp.append(y)\n",
    "        return sum(res + i >= k for i in has)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        dct = [[] for _ in range(n)]\n",
    "        for i, j in edges:\n",
    "            dct[i].append(j)\n",
    "            dct[j].append(i)\n",
    "\n",
    "        guess = set((a, b) for a, b in guesses)\n",
    "\n",
    "        father = [[0, 0] for _ in range(n)]\n",
    "        stack = [[0, -1]]\n",
    "        total_t = 0\n",
    "        while stack:\n",
    "            i, fa = stack.pop()\n",
    "            for j in dct[i]:\n",
    "                if j != fa:\n",
    "                    cur_pre, cur_t = father[i][:]\n",
    "                    if (i, j) in guess:\n",
    "                        cur_pre += 1\n",
    "                        cur_t += 1\n",
    "                        total_t += 1\n",
    "                    if (j, i) in guess:\n",
    "                        cur_pre += 1\n",
    "                    father[j] = [cur_pre, cur_t]\n",
    "                    stack.append([j, i])\n",
    "\n",
    "        return sum(total_t - t + pre - t  >= k for pre, t in father)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        f: List[int] = [0 for _ in range(len(edges) + 1)]\n",
    "        e: Dict[List[int]] = {k: [] for k in range(len(edges) + 1)}\n",
    "\n",
    "        for each in edges:\n",
    "            e[each[0]].append(each[1])\n",
    "            e[each[1]].append(each[0])\n",
    "\n",
    "        s = {(x, y) for x, y in guesses}\n",
    "        \n",
    "        def bfs(start: int):\n",
    "            queue: List[int] = []\n",
    "            queue.append((start, -1))\n",
    "            while len(queue) > 0:\n",
    "                front, pre_front = queue.pop(0)\n",
    "                for des in e[front]:\n",
    "                    if des == pre_front: continue\n",
    "                    queue.append((des, front))\n",
    "                    if (front, des) in s:\n",
    "                        f[start] += 1\n",
    "\n",
    "        def bfs_(start: int):\n",
    "            queue: List[int] = []\n",
    "            queue.append((start, -1))\n",
    "            while len(queue) > 0:\n",
    "                front, pre_front = queue.pop(0)\n",
    "                for des in e[front]:\n",
    "                    if des == pre_front: continue\n",
    "                    queue.append((des, front))\n",
    "                    f[des] = f[front] - ((front, des) in s) + ((des, front) in s)\n",
    "        \n",
    "        bfs(0)\n",
    "        bfs_(0)\n",
    "        result = 0\n",
    "\n",
    "        for each in f:\n",
    "            if each >= k:\n",
    "                result += 1\n",
    "\n",
    "        return result\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        dic = [set() for _ in range(n)]\n",
    "        for x, y in guesses:\n",
    "            dic[y].add(x)\n",
    "        \n",
    "        \n",
    "        cur = len(guesses) - len(dic[0])\n",
    "        q = [0]\n",
    "        mem = set([0])\n",
    "        while len(q) > 0:\n",
    "            nq = []\n",
    "            for i in q:\n",
    "                for j in g[i]:\n",
    "                    if j in mem:\n",
    "                        continue\n",
    "                    mem.add(j)\n",
    "                    cur -= len(dic[j])\n",
    "                    if i in dic[j]:\n",
    "                        cur += 1\n",
    "                    nq.append(j)\n",
    "            q = nq\n",
    "\n",
    "        res = 0\n",
    "        q = [(0, cur)]\n",
    "        if cur >= k:\n",
    "            res += 1\n",
    "        mem = set([0])\n",
    "        while len(q) > 0:\n",
    "            nq = []\n",
    "            for i, prev in q:\n",
    "                for j in g[i]:\n",
    "                    cur = prev\n",
    "                    if j in mem:\n",
    "                        continue\n",
    "                    mem.add(j)\n",
    "                    if i in dic[j]:\n",
    "                        cur -= 1\n",
    "                    if j in dic[i]:\n",
    "                        cur += 1\n",
    "                    nq.append((j, cur))\n",
    "                    if cur >= k:\n",
    "                        res += 1\n",
    "            q = nq\n",
    "        return res\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        #n=len(edges)+1\n",
    "        e=collections.defaultdict(list)\n",
    "        for i,j in edges:\n",
    "            e[i].append(j)\n",
    "            e[j].append(i)\n",
    "        ls=set()\n",
    "        for x,y in guesses:\n",
    "            ls.add((x,y))\n",
    "        guesses=ls\n",
    "        def bfs(x):\n",
    "            q=collections.deque()\n",
    "            q.append([x,-1])\n",
    "            s=set()\n",
    "            while q:\n",
    "                tp,pre=q.popleft()\n",
    "                for i in e[tp]:\n",
    "                    if pre!=i:\n",
    "                        s.add((tp,i))\n",
    "                        q.append([i,tp])\n",
    "            ret=0\n",
    "            for a in guesses:\n",
    "                if a in s:\n",
    "                    ret+=1\n",
    "            return ret\n",
    "        ans=0\n",
    "        now=bfs(0)\n",
    "        q=collections.deque()\n",
    "        q.append([0,-1,now])\n",
    "        while q:\n",
    "            tp,pre,tmp=q.popleft()\n",
    "            if (pre,tp)in guesses:\n",
    "                tmp-=1\n",
    "            if (tp,pre)in guesses:\n",
    "                tmp+=1\n",
    "            #print(tp,pre,now)\n",
    "            if tmp>=k:\n",
    "                ans+=1\n",
    "            for i in e[tp]:\n",
    "                if i!=pre:\n",
    "                    q.append([i,tp,tmp])\n",
    " \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from string import *\n",
    "from re import *\n",
    "from datetime import *\n",
    "from collections import *\n",
    "from heapq import *\n",
    "from bisect import *\n",
    "from copy import *\n",
    "from math import *\n",
    "from random import *\n",
    "from statistics import *\n",
    "from itertools import *\n",
    "from functools import *\n",
    "from operator import *\n",
    "from io import *\n",
    "from sys import *\n",
    "from json import *\n",
    "from builtins import *\n",
    "from typing import *\n",
    "MOD = 1000000007 # 998244353\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "\n",
    "        corr = [set() for _ in range(n)]\n",
    "        for u, v in guesses:\n",
    "            corr[u].add(v)\n",
    "\n",
    "        # step1\n",
    "        root = 0\n",
    "        parent = [-1] * n  # 记录父亲结点\n",
    "        order = [root]  # 记录dfs序，深度单调增加\n",
    "        stack = [root]\n",
    "        while stack:\n",
    "            from_node = stack.pop()\n",
    "            for to_node in g[from_node]:\n",
    "                if to_node == parent[from_node]:\n",
    "                    continue\n",
    "                parent[to_node] = from_node\n",
    "                order.append(to_node)\n",
    "                stack.append(to_node)\n",
    "        \n",
    "        # step2\n",
    "        e = lambda _: 0\n",
    "        merge = add\n",
    "        def op(to_val, from_node, to_node):\n",
    "            return to_val + int(to_node in corr[from_node])\n",
    "        \n",
    "        dp1 = list(map(e, range(n)))\n",
    "        dp2 = list(map(e, range(n)))\n",
    "        \n",
    "        for from_node in order[::-1]:\n",
    "            t = e(from_node)\n",
    "            for to_node in g[from_node]:\n",
    "                if parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = t\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            t = e(from_node)\n",
    "            for to_node in g[from_node][::-1]:\n",
    "                if parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = merge(t, dp2[to_node])\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            dp1[from_node] = t\n",
    "\n",
    "        # step3\n",
    "        for new_root in order[1:]:  # 把最终结果放入dp1，res[u] = merge(dp1[u], dp2[u])\n",
    "            par = parent[new_root]\n",
    "            dp2[new_root] = op(merge(dp2[new_root], dp2[par]), new_root, par)\n",
    "            # if new_root in corr[par]:\n",
    "            #     dp2[new_root] -= 1\n",
    "            dp1[new_root] = merge(dp1[new_root], dp2[new_root])\n",
    "        \n",
    "        # print(dp1)\n",
    "        return sum(x >= k for x in dp1)\n",
    "\n",
    "# testcases = [\n",
    "#     [[[0,1],[1,2],[1,3],[4,2]], [[1,3],[0,1],[1,0],[2,4]], 3],\n",
    "#     # [],\n",
    "# ]\n",
    "\n",
    "# s = Solution()\n",
    "# func_name = dir(s)[-1]\n",
    "# func = getattr(s, func_name)\n",
    "\n",
    "# for args in testcases:\n",
    "#     print(func(*args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        N = len(edges) + 1\n",
    "        d = [set() for _ in range(N)]\n",
    "        for a, b in edges:\n",
    "            d[a].add(b)\n",
    "            d[b].add(a)\n",
    "\n",
    "        gs = {(u, v) for u, v in guesses}\n",
    "        all_ok_num = 0\n",
    "        ps = [0] * N\n",
    "        ps[0] = (0, 0)\n",
    "\n",
    "        dp = {0}\n",
    "        while 1:\n",
    "            now = set()\n",
    "            for a in dp:\n",
    "                ok_num, ng_num = ps[a]\n",
    "\n",
    "                for b in d[a]:\n",
    "                    if ps[b]:\n",
    "                        continue\n",
    "\n",
    "                    now.add(b)\n",
    "                    if (a, b) in gs:\n",
    "                        all_ok_num += 1\n",
    "                        _ok_num = ok_num + 1\n",
    "                    else:\n",
    "                        _ok_num = ok_num\n",
    "\n",
    "                    ps[b] = _ok_num, ng_num + 1 if (b, a) in gs else ng_num\n",
    "\n",
    "            if not now:\n",
    "                break\n",
    "            dp = now\n",
    "\n",
    "        ret = 0\n",
    "        for ok_num, ng_num in ps:\n",
    "            if all_ok_num - ok_num + ng_num >= k:\n",
    "                ret += 1\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Rerooting:\n",
    "    def __init__(self, n: int, decrement: int = 0):\n",
    "\n",
    "        self.n = n\n",
    "        self.adj = [[] for _ in range(n)]\n",
    "        self.root = None\n",
    "        self.decrement = decrement\n",
    "\n",
    "    def add_edge(self, u: int, v: int):\n",
    "\n",
    "        u -= self.decrement\n",
    "        v -= self.decrement\n",
    "        self.adj[u].append(v)\n",
    "        self.adj[v].append(u)\n",
    "\n",
    "    def rerooting(self, op, merge, e, root: int = 0) -> list:\n",
    "\n",
    "        root -= self.decrement\n",
    "        assert 0 <= root < self.n\n",
    "        self.root = root\n",
    "        self.parent = [-1] * self.n\n",
    "        self.order = [root]\n",
    "        stack = [root]\n",
    "        while stack:\n",
    "            from_node = stack.pop()\n",
    "            for to_node in self.adj[from_node]:\n",
    "                if to_node == self.parent[from_node]:\n",
    "                    continue\n",
    "                self.parent[to_node] = from_node\n",
    "                self.order.append(to_node)\n",
    "                stack.append(to_node)\n",
    "\n",
    "        dp1 = [e] * self.n\n",
    "        dp2 = [e] * self.n\n",
    "        for from_node in self.order[::-1]:\n",
    "            t = e\n",
    "            for to_node in self.adj[from_node]:\n",
    "                if self.parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = t\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            t = e\n",
    "            for to_node in self.adj[from_node][::-1]:\n",
    "                if self.parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = merge(t, dp2[to_node])\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            dp1[from_node] = t\n",
    "\n",
    "        for new_root in self.order[1:]:\n",
    "            par = self.parent[new_root]\n",
    "            dp2[new_root] = op(merge(dp2[new_root], dp2[par]), new_root, par)\n",
    "            dp1[new_root] = merge(dp1[new_root], dp2[new_root])\n",
    "        return dp1\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(\n",
    "        self, edges: List[List[int]], guesses: List[List[int]], k: int\n",
    "    ) -> int:\n",
    "        n = len(edges) + 1\n",
    "        T = Rerooting(n, 0)\n",
    "        for u, v in edges:\n",
    "            T.add_edge(u, v)\n",
    "        g = [set() for _ in range(n)]\n",
    "        for u, v in guesses:\n",
    "            g[u].add(v)\n",
    "\n",
    "        def op(dp_v, u, v):\n",
    "\n",
    "            return dp_v + (v in g[u])\n",
    "\n",
    "        def merge(dp_pu, dp_pv):\n",
    "\n",
    "            return dp_pu + dp_pv\n",
    "\n",
    "        e = 0\n",
    "        dp = T.rerooting(op=op, merge=merge, e=e, root=0)\n",
    "        return sum(i >= k for i in dp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        tree={}\n",
    "        for i in edges:\n",
    "            if i[0] in tree:\n",
    "                tree[i[0]].append(i[1])\n",
    "            else:\n",
    "                tree[i[0]]=[i[1]]\n",
    "            if i[1] in tree:\n",
    "                tree[i[1]].append(i[0])\n",
    "            else:\n",
    "                tree[i[1]]=[i[0]]\n",
    "        \n",
    "        start_node=list(tree.keys())[0]\n",
    "        treeset=set()\n",
    "        nodelist=[[None,start_node]]\n",
    "        while True:\n",
    "            if len(nodelist)==0:\n",
    "                break\n",
    "            tmpnodelist=[]\n",
    "            for i in nodelist:\n",
    "                for j in tree[i[1]]:\n",
    "                    if j!=i[0]:\n",
    "                        treeset.add((i[1],j))\n",
    "                        tmpnodelist.append([i[1],j])\n",
    "            nodelist=tmpnodelist\n",
    "        \n",
    "        result={}\n",
    "        gset=set([(i[0],i[1]) for i in guesses])\n",
    "        tmpresult=0\n",
    "        for i in gset:\n",
    "            if i in treeset:\n",
    "                tmpresult+=1\n",
    "        result[start_node]=tmpresult\n",
    "\n",
    "        nodelist=[[None,start_node]]\n",
    "        while True:\n",
    "            if len(nodelist)==0:\n",
    "                break\n",
    "            tmpnodelist=[]\n",
    "            for i in nodelist:\n",
    "                for j in tree[i[1]]:\n",
    "                    if j!=i[0]:\n",
    "                        tmpcnt=0\n",
    "                        if (i[1],j) in gset:\n",
    "                            tmpcnt-=1\n",
    "                        if (j,i[1]) in gset:\n",
    "                            tmpcnt+=1\n",
    "                        result[j]=result[i[1]]+tmpcnt\n",
    "                        tmpnodelist.append([i[1],j])\n",
    "            nodelist=tmpnodelist\n",
    "        \n",
    "        cnt=0\n",
    "        for i in result:\n",
    "            if result[i]>=k:\n",
    "                cnt+=1\n",
    "        return cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        tree = defaultdict(set)\n",
    "        for l,r in edges:\n",
    "            tree[l].add(r)\n",
    "            tree[r].add(l)\n",
    "            \n",
    "        BFS = [0]\n",
    "        for p in BFS:\n",
    "            for c in tree[p]:\n",
    "                tree[c].remove(p)\n",
    "                if tree[c]:\n",
    "                    BFS.append(c)\n",
    "        \n",
    "        guesscnt = Counter(map(tuple, guesses))\n",
    "                \n",
    "        score = {0: sum(c in tree[p] for p,c in guesses)}\n",
    "        for p in BFS:\n",
    "            scorep = score[p]\n",
    "            for c in tree[p]:\n",
    "                score[c] = scorep - guesscnt[p,c] + guesscnt[c,p]\n",
    "        return sum(i>=k for i in score.values())\n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        g=dict()\n",
    "        for i in guesses:\n",
    "            g[tuple(i)]=g.get(tuple(i),0)+1\n",
    "        dic=dict()\n",
    "        for i in edges:\n",
    "            s=dic.get(i[0],set())\n",
    "            s.add(i[1])\n",
    "            dic[i[0]]=s\n",
    "            s=dic.get(i[1],set())\n",
    "            s.add(i[0])\n",
    "            dic[i[1]]=s\n",
    "        count=0\n",
    "        m=max(list(dic.keys()))+1\n",
    "        F=[0]*m\n",
    "        F[0]=-1\n",
    "        next=[0]\n",
    "        while next!=[]:\n",
    "            now=next[:]\n",
    "            next=[]\n",
    "            for i in now:\n",
    "                for j in dic[i]:\n",
    "                    if F[j]==0:\n",
    "                        next.append(j)\n",
    "                        F[j]=1\n",
    "                        count+=g.get((i,j),0)\n",
    "        memo={0:count}\n",
    "        F=[0]*m\n",
    "        F[0]=-1\n",
    "        next=[0]\n",
    "        while next!=[]:\n",
    "            now=next[:]\n",
    "            next=[]\n",
    "            for i in now:\n",
    "                for j in dic[i]:\n",
    "                    if F[j]==0:\n",
    "                        next.append(j)\n",
    "                        F[j]=1\n",
    "                        count=memo[i]\n",
    "                        count-=g.get((i,j),0)\n",
    "                        count+=g.get((j,i),0)\n",
    "                        memo[j]=count\n",
    "        ans=0\n",
    "        for i in memo:\n",
    "            if memo[i]>=k:\n",
    "                ans+=1\n",
    "        return ans\n",
    "            \n",
    "                            \n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from string import *\n",
    "from re import *\n",
    "from datetime import *\n",
    "from collections import *\n",
    "from heapq import *\n",
    "from bisect import *\n",
    "from copy import *\n",
    "from math import *\n",
    "from random import *\n",
    "from statistics import *\n",
    "from itertools import *\n",
    "from functools import *\n",
    "from operator import *\n",
    "from io import *\n",
    "from sys import *\n",
    "from json import *\n",
    "from builtins import *\n",
    "from typing import *\n",
    "MOD = 1000000007 # 998244353\n",
    "\n",
    "class Rerooting():\n",
    "    \"\"\"https://atcoder.jp/contests/dp/submissions/22766939\"\"\"\n",
    "    __slots__ = (\"adjList\", \"_n\", \"_decrement\", \"_root\", \"_parent\", \"_order\")\n",
    "\n",
    "    def __init__(self, n: int, decrement: int = 1):\n",
    "        \"\"\"n是顶点数， decrement是1-index 或者 0-indexed，默认1-index\"\"\"\n",
    "        self._n = n\n",
    "        self.adjList = [[] for _ in range(n)]\n",
    "        self._root = None  # 最初遍历时的根\n",
    "        self._decrement = decrement\n",
    "\n",
    "    def add_edge(self, u: int, v: int):\n",
    "        \"\"\"添加一条无向边\"\"\"\n",
    "        u -= self._decrement\n",
    "        v -= self._decrement\n",
    "        self.adjList[u].append(v)\n",
    "        self.adjList[v].append(u)\n",
    "\n",
    "    def rerooting(self, op, merge, e, root=-1):\n",
    "        \"\"\"\n",
    "        <概要>\n",
    "        1. 以root为根，先用dfs求出一次树结构，计算出树的dfs序：self.order\n",
    "        2. dp1[u]记录以u为根的子树值，dp2[u]记录u往father指的子树的值\n",
    "           由于树结构固定，dp1可以自下而上推算，dp2暂存father指向的除了自己外的子树的值\n",
    "        3. 换根时的dp： \n",
    "            dp2[u]还缺dp2[father]的信息，所以：dp2[u] = op(merge(dp2[u], dp2[fa]), u, fa)\n",
    "            dp1[u] = res = merge(dp1[u], dp2[u])\n",
    "        复杂度 O(|V|) (V是树的顶点数)\n",
    "        参考： https://qiita.com/keymoon/items/2a52f1b0fb7ef67fb89e\n",
    "        \"\"\"\n",
    "        # step1\n",
    "        if root == -1:\n",
    "            root = 0\n",
    "        else:\n",
    "            root -= self._decrement\n",
    "        assert 0 <= root < self._n\n",
    "        self._parent = parent = [-1] * self._n  # 记录父亲结点\n",
    "        self._root = order = [root]  # 记录dfs序，深度单调增加\n",
    "        stack = [root]\n",
    "        while stack:\n",
    "            from_node = stack.pop()\n",
    "            for to_node in self.adjList[from_node]:\n",
    "                if to_node == parent[from_node]:\n",
    "                    continue\n",
    "                parent[to_node] = from_node\n",
    "                order.append(to_node)\n",
    "                stack.append(to_node)\n",
    "        # step2\n",
    "        dp1 = list(map(e, range(self._n)))\n",
    "        dp2 = list(e(i) for i in range(self._n))\n",
    "        for from_node in order[::-1]:\n",
    "            t = e(from_node)\n",
    "            for to_node in self.adjList[from_node]:\n",
    "                if parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = t\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            t = e(from_node)\n",
    "            for to_node in self.adjList[from_node][::-1]:\n",
    "                if parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = merge(t, dp2[to_node])\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            dp1[from_node] = t\n",
    "        # step3\n",
    "        for new_root in order[1:]:  # 把最终结果放入dp1，res[u] = merge(dp1[u], dp2[u])\n",
    "            par = parent[new_root]\n",
    "            dp2[new_root] = op(merge(dp2[new_root], dp2[par]), new_root, par)\n",
    "            dp1[new_root] = merge(dp1[new_root], dp2[new_root])\n",
    "        return dp1\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        rt = Rerooting(len(edges) + 1, 0)\n",
    "        for u, v in edges:\n",
    "            rt.add_edge(u, v)\n",
    "\n",
    "        corr = defaultdict(set)\n",
    "        for u, v in guesses:\n",
    "            corr[u].add(v)\n",
    "        \n",
    "        def op(val, u, v):\n",
    "            return val + (1 if v in corr[u] else 0)\n",
    "        res = rt.rerooting(op, add, lambda _: 0)\n",
    "        \n",
    "        return sum(x >= k for x in res)\n",
    "        \n",
    "\n",
    "# testcases = [\n",
    "#     [[[0,1],[1,2],[1,3],[4,2]], [[1,3],[0,1],[1,0],[2,4]], 3],\n",
    "#     # [],\n",
    "# ]\n",
    "\n",
    "# s = Solution()\n",
    "# func_name = dir(s)[-1]\n",
    "# func = getattr(s, func_name)\n",
    "\n",
    "# for args in testcases:\n",
    "#     print(func(*args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        g = collections.defaultdict(int)\n",
    "        n = len(edges)+1\n",
    "        parent = [-1]*n\n",
    "        parent[0] = -2\n",
    "        dp = [0]*n\n",
    "        con = [[] for _ in range(n)]\n",
    "        children = [[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            con[x].append(y)\n",
    "            con[y].append(x)\n",
    "        queue = [0]\n",
    "        i = 0\n",
    "        while i < len(queue):\n",
    "            for j in con[queue[i]]:\n",
    "                if parent[j] == -1:\n",
    "                    parent[j] = queue[i]\n",
    "                    queue.append(j)\n",
    "                    children[queue[i]].append(j)\n",
    "            i += 1\n",
    "        for x,y in guesses:\n",
    "            g[(x,y)] += 1\n",
    "            if parent[y] == x:\n",
    "                dp[0] += 1\n",
    "        queue = [0]\n",
    "        i = 0\n",
    "        while i < len(queue):\n",
    "            for j in children[queue[i]]:\n",
    "                dp[j] = dp[queue[i]]\n",
    "                dp[j] += g[(j,queue[i])]\n",
    "                dp[j] -= g[(queue[i],j)]\n",
    "                queue.append(j)\n",
    "            i += 1\n",
    "        ans = 0\n",
    "        for i in dp:\n",
    "            if i >= k:\n",
    "                ans += 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        edgeDict = defaultdict(set)\n",
    "        for u, v in edges:\n",
    "            edgeDict[u].add(v)\n",
    "            edgeDict[v].add(u)\n",
    "            \n",
    "        guesses = Counter(map(tuple, guesses))\n",
    "        \n",
    "        path = []\n",
    "        cnt = defaultdict(int)\n",
    "        d = deque([0])\n",
    "        asked = set()\n",
    "        while d:\n",
    "            u = d.popleft()\n",
    "            asked.add(u)\n",
    "            for v in edgeDict[u]:\n",
    "                if v not in asked:\n",
    "                    d.append(v)\n",
    "                    path.append((u, v))\n",
    "                    cnt[0] += guesses[(u, v)]\n",
    "        \n",
    "        for u, v in path:\n",
    "            cnt[v] = cnt[u] + guesses[(v, u)] - guesses[(u, v)]\n",
    "        \n",
    "        return sum(cnt[u] >= k for u in range(len(edges) + 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        #Build the tree\n",
    "     \n",
    "        adj=defaultdict(list)\n",
    "        for a,b in edges:\n",
    "            adj[a].append(b)\n",
    "            adj[b].append(a)\n",
    "\n",
    "        gpdic=defaultdict(set)       \n",
    "        for p,c in guesses:    \n",
    "            gpdic[p].add(c)       \n",
    "        #pick any as root\n",
    "        nodes=list(adj.keys())\n",
    "        children=defaultdict(list)\n",
    "        visit={nodes[0]}\n",
    "        dfs=[nodes[0]] \n",
    "        iright=0\n",
    "        while dfs:\n",
    "            cur=dfs.pop()            \n",
    "            for child in adj[cur]:\n",
    "                if child not in visit:\n",
    "                    visit.add(child)\n",
    "                    if child in gpdic[cur]:\n",
    "                        iright+=1\n",
    "                    children[cur].append(child)\n",
    "                    dfs.append(child)\n",
    "        ans=0\n",
    "       \n",
    "        bfs=[(nodes[0],iright)]\n",
    "        while bfs:\n",
    "            cur,right=bfs.pop()\n",
    "            if right>=k:\n",
    "                ans+=1\n",
    "            for child in children[cur]:\n",
    "                cright=right\n",
    "                if child in gpdic[cur]:\n",
    "                    cright-=1\n",
    "                if cur in gpdic[child]:\n",
    "                    cright+=1\n",
    "                bfs.append((child,cright)) \n",
    "    \n",
    "        return ans\n",
    "                \n",
    "\n",
    "        \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(\n",
    "        self, \n",
    "        edges: List[List[int]], \n",
    "        guesses: List[List[int]], \n",
    "        k: int\n",
    "    ) -> int:\n",
    "        \n",
    "        n = len(edges) + 1\n",
    "        \n",
    "        graph = defaultdict(set)\n",
    "        for u, v in edges:\n",
    "            graph[u].add(v)\n",
    "            graph[v].add(u)\n",
    "            \n",
    "        guess = set((u, v) for u, v in guesses)\n",
    "        \n",
    "        # dp[i]: 表示以结点 i 为根，边猜测正确的个数\n",
    "        dp = [0 for _ in range(n)]\n",
    "\n",
    "        # 第一次 BFS：计算 dp[0]，得到 BFS 生成树\n",
    "        que = deque()\n",
    "        que.append(0)\n",
    "        \n",
    "        tree = defaultdict(set) # BFS 生成树\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            for v in graph[u]:\n",
    "                if v not in tree:\n",
    "                    tree[u].add(v)\n",
    "                    que.append(v)\n",
    "\n",
    "                    if (u, v) in guess:\n",
    "                        dp[0] += 1\n",
    "        \n",
    "        # 第二次 BFS：遍历 `BFS 生成树` 中的每一条边 (u, v)\n",
    "        # 分别以 u 和 v 为根，使用动态规划计算 dp[i]\n",
    "        que.append(0)\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            \n",
    "            for v in tree[u]:\n",
    "                que.append(v)\n",
    "\n",
    "                # 下面是状态转移\n",
    "                if (u, v) in guess:\n",
    "                    # 1. 以 u 为根，边 (u, v) 猜中时\n",
    "                    if (v, u) in guess:\n",
    "                        # 1.1 以 v 为根，边 (v, u) 猜中时\n",
    "                        # 此时边猜中的个数不变\n",
    "                        dp[v] = dp[u]\n",
    "                    else:\n",
    "                        # 1.2 以 v 为根，边 (v, u) 未猜中\n",
    "                        # 此时边猜中的个数少一个，即 (u, v)\n",
    "                        dp[v] = dp[u] - 1\n",
    "                else:\n",
    "                    # 2. 以 u 为根，边 (u, v) 未猜中时\n",
    "                    if (v, u) in guess:\n",
    "                        # 2.1 以 v 为根，边 (v, u) 猜中时\n",
    "                        # 此时边猜中的个数多一个，即 (v, u)\n",
    "                        dp[v] = dp[u] + 1\n",
    "                    else:\n",
    "                        # 2.2 以 v 为根，边 (v, u) 未猜中时\n",
    "                        # 此时边猜中的个数不变\n",
    "                        dp[v] = dp[u]\n",
    "                            \n",
    "        return sum(cnt >= k for cnt in dp)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(\n",
    "        self, \n",
    "        edges: List[List[int]], \n",
    "        guesses: List[List[int]], \n",
    "        k: int\n",
    "    ) -> int:\n",
    "        \n",
    "        n = len(edges) + 1\n",
    "        \n",
    "        graph = defaultdict(set)\n",
    "        for u, v in edges:\n",
    "            graph[u].add(v)\n",
    "            graph[v].add(u)\n",
    "            \n",
    "        guess = set((u, v) for u, v in guesses)\n",
    "        \n",
    "        # dp[i]: 表示以结点 i 为根，边猜测正确的个数\n",
    "        dp = [0 for _ in range(n)]\n",
    "\n",
    "        # 第一次 BFS：计算 dp[0]，得到 BFS 生成树\n",
    "        que = deque()\n",
    "        que.append(0)\n",
    "        \n",
    "        tree = defaultdict(set) # BFS 生成树\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            for v in graph[u]:\n",
    "                if v not in tree:\n",
    "                    tree[u].add(v)\n",
    "                    que.append(v)\n",
    "\n",
    "                    if (u, v) in guess:\n",
    "                        dp[0] += 1\n",
    "        \n",
    "        # 第二次 BFS：遍历 `BFS 生成树` 中的每一条边 (u, v)\n",
    "        # 分别以 u 和 v 为根，使用动态规划计算 dp[i]\n",
    "        que.append(0)\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            \n",
    "            for v in tree[u]:\n",
    "                que.append(v)\n",
    "\n",
    "                # 下面是状态转移\n",
    "                if (u, v) in guess:\n",
    "                    # 1. 以 u 为根，边 (u, v) 猜中时\n",
    "                    if (v, u) in guess:\n",
    "                        # 1.1 以 v 为根，边 (v, u) 猜中时\n",
    "                        # 此时边猜中的个数不变\n",
    "                        dp[v] = dp[u]\n",
    "                    else:\n",
    "                        # 1.2 以 v 为根，边 (v, u) 未猜中\n",
    "                        # 此时边猜中的个数少一个，即 (u, v)\n",
    "                        dp[v] = dp[u] - 1\n",
    "                else:\n",
    "                    # 2. 以 u 为根，边 (u, v) 未猜中时\n",
    "                    if (v, u) in guess:\n",
    "                        # 2.1 以 v 为根，边 (v, u) 猜中时\n",
    "                        # 此时边猜中的个数多一个，即 (v, u)\n",
    "                        dp[v] = dp[u] + 1\n",
    "                    else:\n",
    "                        # 2.2 以 v 为根，边 (v, u) 未猜中时\n",
    "                        # 此时边猜中的个数不变\n",
    "                        dp[v] = dp[u]\n",
    "                            \n",
    "        return sum(cnt >= k for cnt in dp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = max(itertools.chain(*edges)) + 1\n",
    "        d = collections.Counter()\n",
    "        for a, b in edges:\n",
    "            d[a] += 1\n",
    "            d[b] += 1\n",
    "        adj = collections.defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            adj[a].append(b)\n",
    "            adj[b].append(a)\n",
    "        edges = [(a, b) for a, b in edges]\n",
    "        gs = {(a, b) for a, b in guesses}\n",
    "        ones = []\n",
    "        for v in range(n):\n",
    "            if d[v] == 1:\n",
    "                ones.append(v)\n",
    "        cnt = collections.Counter()\n",
    "        st = []\n",
    "        while ones:\n",
    "            a = ones.pop()\n",
    "            for b in adj[a]:\n",
    "                if (b, a) not in cnt:\n",
    "                    cnt[(a, b)] = cnt[a]\n",
    "                    st.append((a, b))\n",
    "                    cnt[b] += cnt[a]\n",
    "                    if (b, a) in gs:\n",
    "                        cnt[(a, b)] += 1\n",
    "                        cnt[b] += 1\n",
    "                    d[b] -= 1\n",
    "                    if d[b] == 1:\n",
    "                        ones.append(b)\n",
    "        assert(len(st) == n - 1)\n",
    "        while st:\n",
    "            a, b = st.pop()\n",
    "            cnt[a] += cnt[b] - cnt[(a, b)]\n",
    "            if (a, b) in gs:\n",
    "                cnt[a] += 1\n",
    "        return sum(cnt[v] >= k for v in cnt if isinstance(v, int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(\n",
    "        self, \n",
    "        edges: List[List[int]], \n",
    "        guesses: List[List[int]], \n",
    "        k: int\n",
    "    ) -> int:\n",
    "        \n",
    "        n = len(edges) + 1\n",
    "        \n",
    "        graph = defaultdict(set)\n",
    "        for u, v in edges:\n",
    "            graph[u].add(v)\n",
    "            graph[v].add(u)\n",
    "            \n",
    "        guess = set((u, v) for u, v in guesses)\n",
    "        \n",
    "        # dp[i]: 表示以结点 i 为根，边猜测正确的个数\n",
    "        dp = [0 for _ in range(n)]\n",
    "\n",
    "        # 第一次 BFS：计算 dp[0]\n",
    "        que = deque()\n",
    "        que.append(0)\n",
    "        \n",
    "        visited = set()\n",
    "        visited.add(0)\n",
    "\n",
    "        tree = defaultdict(set) # BFS 生成树\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            for v in graph[u]:\n",
    "                if v not in visited:\n",
    "                    visited.add(v)\n",
    "                    que.append(v)\n",
    "\n",
    "                    tree[u].add(v)\n",
    "                    \n",
    "                    if (u, v) in guess:\n",
    "                        dp[0] += 1\n",
    "        \n",
    "        # 第二次 BFS：遍历 `BFS 生成树` 中的每一条边 (u, v)\n",
    "        # 分别以 u 和 v 为根，使用动态规划计算 dp[i]\n",
    "        que.append(0)\n",
    "\n",
    "        while que:\n",
    "            u = que.popleft()\n",
    "            \n",
    "            for v in tree[u]:\n",
    "                que.append(v)\n",
    "\n",
    "                if (u, v) in guess:\n",
    "                    if (v, u) in guess:\n",
    "                        dp[v] = dp[u]\n",
    "                    else:\n",
    "                        dp[v] = dp[u] - 1\n",
    "                else:\n",
    "                    if (v, u) in guess:\n",
    "                        dp[v] = dp[u] + 1\n",
    "                    else:\n",
    "                        dp[v] = dp[u]\n",
    "                            \n",
    "        return sum(cnt >= k for cnt in dp)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edg: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        edges={}\n",
    "        n=len(edg)+1\n",
    "        for i in range(n):\n",
    "            edges[i]=set()\n",
    "        for i in edg:\n",
    "            edges[i[0]].add(i[1])\n",
    "            edges[i[1]].add(i[0])\n",
    "        #print(edges)\n",
    "        target=set([(i[0],i[1]) for i in guesses])\n",
    "        #print(target)\n",
    "        tree={}\n",
    "        for i in range(n):\n",
    "            tree[i]=set()\n",
    "        queue=[0]\n",
    "        covered={0}\n",
    "        tmp=[]\n",
    "        while queue:\n",
    "            tmp=[]\n",
    "            for node in queue:\n",
    "                for i in edges[node]:\n",
    "                    if i not in covered:\n",
    "                        tree[node].add(i)\n",
    "                        covered.add(i)\n",
    "                        tmp.append(i)\n",
    "\n",
    "            queue=tmp\n",
    "        #print(tree)\n",
    "        ans=0\n",
    "        cur=0\n",
    "        for i in guesses:\n",
    "            if i[1] in tree[i[0]]:\n",
    "                cur+=1\n",
    "        if cur>=k:\n",
    "            ans+=1\n",
    "        #print(cur)\n",
    "        dic={}\n",
    "        queue=[0]\n",
    "        covered={0}\n",
    "        dic[0]=cur\n",
    "        while queue:\n",
    "            tmp=[]\n",
    "            for node in queue:\n",
    "                \n",
    "                for i in tree[node]:\n",
    "                    if i not in covered:\n",
    "                        flag=0\n",
    "                        if (node,i) in target:\n",
    "                            flag-=1\n",
    "                        if (i,node) in target:\n",
    "                            flag+=1\n",
    "                        dic[i]=dic[node]+flag\n",
    "                        covered.add(i)\n",
    "                        tmp.append(i)\n",
    "            queue=tmp\n",
    "        ans=0\n",
    "        for i in dic:\n",
    "            if dic[i]>=k:\n",
    "                ans+=1\n",
    "        return ans\n",
    "            \n",
    "            \n",
    "                        \n",
    "            \n",
    "            \n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: list[list[int]], guesses: list[list[int]], k: int) -> int:\n",
    "        n = 1 + len(edges)\n",
    "        \n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "\n",
    "        vis = set()\n",
    "        for u, v in guesses:\n",
    "            vis.add((u, v))\n",
    "\n",
    "        cnt = [0] * n\n",
    "\n",
    "        def dfs(x, fa, c):\n",
    "            if (x, fa) in vis:\n",
    "                c += 1\n",
    "                vis.remove((x, fa))\n",
    "            if (fa, x) in vis:\n",
    "                nonlocal k\n",
    "                k -= 1\n",
    "                c -= 1\n",
    "                vis.remove((fa, x))\n",
    "            cnt[x] += c\n",
    "            for y in g[x]:\n",
    "                if y == fa: continue\n",
    "                dfs(y, x, c)\n",
    "\n",
    "        while vis:\n",
    "            (a, b) = vis.pop()\n",
    "            dfs(a, b, 1)\n",
    "\n",
    "        return sum(x >= k for x in cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from types import GeneratorType\n",
    "def bootstrap(f, stack=[]):\n",
    "    def wrappedfunc(*args, **kwargs):\n",
    "        if stack:\n",
    "            return f(*args, **kwargs)\n",
    "        else:\n",
    "            to = f(*args, **kwargs)\n",
    "            while True:\n",
    "                if type(to) is GeneratorType:\n",
    "                    stack.append(to)\n",
    "                    to = next(to)\n",
    "                else:\n",
    "                    stack.pop()\n",
    "                    if not stack:\n",
    "                        break\n",
    "                    to = stack[-1].send(to)\n",
    "            return to\n",
    "    return wrappedfunc\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        d = [[] for _ in range(n)]\n",
    "        for i, j in edges:\n",
    "            d[i].append(j)\n",
    "            d[j].append(i)\n",
    "        g = set((i, j) for i, j in guesses)\n",
    "        c = [0] * n\n",
    "        rc = [0] * n\n",
    "        @bootstrap\n",
    "        def dfs1(i, fa):\n",
    "            for j in d[i]:\n",
    "                if j == fa:\n",
    "                    continue\n",
    "                yield dfs1(j, i)\n",
    "                c[i] += c[j]\n",
    "                if (i, j) in g:\n",
    "                    c[i] += 1\n",
    "            yield 0\n",
    "        dfs1(0, -1)\n",
    "        @bootstrap\n",
    "        def dfs2(i, fa):\n",
    "            for j in d[i]:\n",
    "                if j == fa:\n",
    "                    continue\n",
    "                if (j, i) in g:\n",
    "                    rc[j] += 1\n",
    "                rc[j] += rc[i] + c[i] - c[j] - int((i, j) in g)\n",
    "                yield dfs2(j, i)\n",
    "            yield 0\n",
    "        dfs2(0, -1)\n",
    "        ans = 0\n",
    "        for i, j in zip(c, rc):\n",
    "            if i + j >= k:\n",
    "                ans += 1\n",
    "        # for i in rc:\n",
    "        #     if i >= k:\n",
    "        #         ans += 1\n",
    "        return ans\n",
    "                \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: list[list[int]], guesses: list[list[int]], k: int) -> int:\n",
    "        n = 1 + len(edges)\n",
    "        \n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "\n",
    "        vis = set()\n",
    "        for u, v in guesses:\n",
    "            vis.add((u, v))\n",
    "\n",
    "        cnt = [0] * n\n",
    "\n",
    "        def dfs(x, fa, c):\n",
    "            if (x, fa) in vis:\n",
    "                c += 1\n",
    "                vis.remove((x, fa))\n",
    "            if (fa, x) in vis:\n",
    "                nonlocal k\n",
    "                k -= 1\n",
    "                c -= 1\n",
    "                vis.remove((fa, x))\n",
    "            cnt[x] += c\n",
    "            for y in g[x]:\n",
    "                if y == fa: continue\n",
    "                dfs(y, x, c)\n",
    "\n",
    "        for a, b in vis:\n",
    "            break\n",
    "        while vis:\n",
    "            dfs(b, -1, 0)\n",
    "            for a, b in vis:\n",
    "                break\n",
    "\n",
    "        return sum(x >= k for x in cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from itertools import accumulate\n",
    "from typing import List\n",
    "\n",
    "from typing import List, Tuple\n",
    "\n",
    "\n",
    "class DFSOrder:\n",
    "    __slots__ = (\"starts\", \"ends\", \"_n\", \"_tree\", \"_dfsId\")\n",
    "\n",
    "    def __init__(self, n: int, tree: List[List[int]], root=0) -> None:\n",
    "        \"\"\"dfs序\n",
    "\n",
    "        Args:\n",
    "            n (int): 树节点从0开始,根节点为0\n",
    "            tree (Tree): 无向图邻接表\n",
    "\n",
    "        1. 按照dfs序遍历k个结点形成的回路 每条边恰好经过两次\n",
    "        \"\"\"\n",
    "        self.starts = [0] * n\n",
    "        self.ends = [0] * n\n",
    "        self._n = n\n",
    "        self._tree = tree\n",
    "        self._dfsId = 0\n",
    "        self._dfs(root, -1)\n",
    "\n",
    "    def querySub(self, root: int) -> Tuple[int, int]:\n",
    "        \"\"\"求子树映射到的区间\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "        Returns:\n",
    "            Tuple[int, int]: [start, end] 0 <= start < end <= n\n",
    "        \"\"\"\n",
    "        return self.starts[root], self.ends[root] + 1\n",
    "\n",
    "    def queryId(self, root: int) -> int:\n",
    "        \"\"\"求root自身的dfsId\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "        Returns:\n",
    "            int: id  1 <= id <= n\n",
    "        \"\"\"\n",
    "        return self.ends[root]\n",
    "\n",
    "    def isAncestor(self, root: int, child: int) -> bool:\n",
    "        \"\"\"判断root是否是child的祖先\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "            child (int): 子节点\n",
    "\n",
    "        应用:枚举边时给树的边定向\n",
    "        ```\n",
    "        if not D.isAncestor(e[0], e[1]):\n",
    "            e[0], e[1] = e[1], e[0]\n",
    "        ```\n",
    "        \"\"\"\n",
    "        left1, right1 = self.starts[root], self.ends[root]\n",
    "        left2, right2 = self.starts[child], self.ends[child]\n",
    "        return left1 <= left2 <= right2 <= right1\n",
    "\n",
    "    def _dfs(self, cur: int, pre: int) -> None:\n",
    "        self.starts[cur] = self._dfsId\n",
    "        for next in self._tree[cur]:\n",
    "            if next == pre:\n",
    "                continue\n",
    "            self._dfs(next, cur)\n",
    "        self.ends[cur] = self._dfsId\n",
    "        self._dfsId += 1\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        \"\"\"对每对查询,看能对哪些区间的根节点产生贡献\"\"\"\n",
    "        n = len(edges) + 1\n",
    "        adjList = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            adjList[u].append(v)\n",
    "            adjList[v].append(u)\n",
    "        D = DFSOrder(n, adjList, root=0)\n",
    "\n",
    "        diff = [0] * (n + 1)  # dfs序对应的根节点的合法数\n",
    "        for u, v in guesses:\n",
    "            if D.isAncestor(u, v):\n",
    "                start, end = D.querySub(v)  # v的子树之外的节点可以作为根节点(不包括v)\n",
    "                diff[0] += 1\n",
    "                diff[start] -= 1\n",
    "                diff[end] += 1\n",
    "                diff[n] -= 1\n",
    "            else:\n",
    "                start, end = D.querySub(u)  # u的子树内的节点可以作为根节点\n",
    "                diff[start] += 1\n",
    "                diff[end] -= 1\n",
    "\n",
    "        diff = list(accumulate(diff))\n",
    "        diff = diff[:n]\n",
    "        return sum(x >= k for x in diff)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List, Tuple\n",
    "\n",
    "\n",
    "class DFSOrder:\n",
    "    __slots__ = (\"starts\", \"ends\", \"_n\", \"_tree\", \"_dfsId\")\n",
    "\n",
    "    def __init__(self, n: int, tree: List[List[int]], root=0) -> None:\n",
    "        \"\"\"dfs序\n",
    "\n",
    "        Args:\n",
    "            n (int): 树节点从0开始,根节点为0\n",
    "            tree (Tree): 无向图邻接表\n",
    "\n",
    "        1. 按照dfs序遍历k个结点形成的回路 每条边恰好经过两次\n",
    "        \"\"\"\n",
    "        self.starts = [0] * n\n",
    "        self.ends = [0] * n\n",
    "        self._n = n\n",
    "        self._tree = tree\n",
    "        self._dfsId = 1\n",
    "        self._dfs(root, -1)\n",
    "\n",
    "    def querySub(self, root: int) -> Tuple[int, int]:\n",
    "        \"\"\"求子树映射到的区间\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "        Returns:\n",
    "            Tuple[int, int]: [start, end] 1 <= start <= end <= n\n",
    "        \"\"\"\n",
    "        return self.starts[root], self.ends[root]\n",
    "\n",
    "    def queryId(self, root: int) -> int:\n",
    "        \"\"\"求root自身的dfsId\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "        Returns:\n",
    "            int: id  1 <= id <= n\n",
    "        \"\"\"\n",
    "        return self.ends[root]\n",
    "\n",
    "    def isAncestor(self, root: int, child: int) -> bool:\n",
    "        \"\"\"判断root是否是child的祖先\n",
    "\n",
    "        Args:\n",
    "            root (int): 根节点\n",
    "            child (int): 子节点\n",
    "\n",
    "        应用:枚举边时给树的边定向\n",
    "        ```\n",
    "        if not D.isAncestor(e[0], e[1]):\n",
    "            e[0], e[1] = e[1], e[0]\n",
    "        ```\n",
    "        \"\"\"\n",
    "        left1, right1 = self.starts[root], self.ends[root]\n",
    "        left2, right2 = self.starts[child], self.ends[child]\n",
    "        return left1 <= left2 <= right2 <= right1\n",
    "\n",
    "    def _dfs(self, cur: int, pre: int) -> None:\n",
    "        self.starts[cur] = self._dfsId\n",
    "        for next in self._tree[cur]:\n",
    "            if next == pre:\n",
    "                continue\n",
    "            self._dfs(next, cur)\n",
    "        self.ends[cur] = self._dfsId\n",
    "        self._dfsId += 1\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        \"\"\"对每对查询,看能对哪些区间的根节点产生贡献\"\"\"\n",
    "        n = len(edges) + 1\n",
    "        adjList = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            adjList[u].append(v)\n",
    "            adjList[v].append(u)\n",
    "        D = DFSOrder(n, adjList, root=0)\n",
    "        diff = [0] * (n + 10)\n",
    "        for u, v in guesses:\n",
    "            if D.isAncestor(u, v):\n",
    "                start, end = D.querySub(v)\n",
    "                diff[1] += 1\n",
    "                diff[start] -= 1\n",
    "                diff[end + 1] += 1\n",
    "            else:\n",
    "                start, end = D.querySub(u)\n",
    "                diff[start] += 1\n",
    "                diff[end + 1] -= 1\n",
    "\n",
    "        diff = list(accumulate(diff))\n",
    "        diff = diff[1 : n + 1]\n",
    "        return sum(x >= k for x in diff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def rootCount(self, edges: list[list[int]], guesses: list[list[int]], k: int) -> int:\n",
    "\n",
    "        n = 1 + len(edges)\n",
    "\n",
    "        \n",
    "\n",
    "        g = [[] for _ in range(n)]\n",
    "\n",
    "        for u, v in edges:\n",
    "\n",
    "            g[u].append(v)\n",
    "\n",
    "            g[v].append(u)\n",
    "\n",
    "        vis = {(u, v) for u, v in guesses}\n",
    "\n",
    "        cnt = [0] * n\n",
    "\n",
    "        def dfs(x, fa, c):\n",
    "\n",
    "            if (x, fa) in vis:\n",
    "\n",
    "                c += 1\n",
    "\n",
    "                vis.remove((x, fa))\n",
    "\n",
    "            if (fa, x) in vis:\n",
    "\n",
    "                nonlocal k\n",
    "\n",
    "                k -= 1\n",
    "\n",
    "                c -= 1\n",
    "\n",
    "                vis.remove((fa, x))\n",
    "\n",
    "            cnt[x] += c\n",
    "\n",
    "            for y in g[x]:\n",
    "\n",
    "                if y == fa: continue\n",
    "\n",
    "                dfs(y, x, c)\n",
    "\n",
    "        dfs(0, -1, 0)\n",
    "\n",
    "        return sum(x >= k for x in cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "\n",
    "\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for i in range(n)]\n",
    "\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        st = set([])\n",
    "\n",
    "        for u, v in guesses:\n",
    "            st.add((u, v))\n",
    "        \n",
    "        \n",
    "\n",
    "        path = []\n",
    "        res = 0\n",
    "\n",
    "        def f(u, fa):\n",
    "            nonlocal res\n",
    "            if (fa, u) in st:\n",
    "                res += 1\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                f(v, u)\n",
    "        \n",
    "        f(0, -1)\n",
    "\n",
    "        ans = 0\n",
    "\n",
    "        def dfs(u, fa):\n",
    "            nonlocal res, ans\n",
    "            \n",
    "            if (fa, u) in st: res -= 1\n",
    "            if (u, fa) in st: res += 1\n",
    "\n",
    "            if res >= k:\n",
    "                ans += 1\n",
    "\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                dfs(v, u)\n",
    "\n",
    "            if (fa, u) in st: res += 1\n",
    "            if (u, fa) in st: res -= 1\n",
    "        \n",
    "        dfs(0, -1)\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        path = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            path[u].append(v)\n",
    "            path[v].append(u)\n",
    "\n",
    "        tmstamp = 0\n",
    "        pos_l = [0] * n\n",
    "        pos_r = [0] * n\n",
    "        def dfs(u, p=-1):\n",
    "            nonlocal tmstamp\n",
    "            pos_l[u] = tmstamp\n",
    "            tmstamp += 1\n",
    "            for v in path[u]:\n",
    "                if v != p:\n",
    "                    dfs(v, u)\n",
    "            pos_r[u] = tmstamp\n",
    "        dfs(0)\n",
    "\n",
    "        diff = [0] * (n + 1)\n",
    "        for u, v in guesses:\n",
    "            if pos_l[u] <= pos_l[v]:\n",
    "                diff[0] += 1\n",
    "                diff[pos_l[v]] -= 1\n",
    "                diff[pos_r[v]] += 1\n",
    "                diff[n] -= 1\n",
    "            else:\n",
    "                diff[pos_l[u]] += 1\n",
    "                diff[pos_r[u]] -= 1\n",
    "\n",
    "        res = list(accumulate(diff))\n",
    "        return sum(r >= k for r in res[:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        q = [[] for _ in range(n)]\n",
    "        for a,b in guesses: q[b].append(a)\n",
    "        \n",
    "        def f(x,pa):\n",
    "            res = int(pa in q[x])\n",
    "            for y in g[x]:\n",
    "                if y == pa: continue\n",
    "                res += f(y,x)\n",
    "            return res\n",
    "        \n",
    "        def dp(x,pa,s):\n",
    "            if s >= k: self.ans += 1\n",
    "            for y in g[x]:\n",
    "                if y == pa: continue\n",
    "                dp(y,x,s - (x in q[y]) + (y in q[x]))\n",
    "        \n",
    "        self.ans = 0\n",
    "        dp(0,-1,f(0,-1))\n",
    "        return self.ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        e = collections.defaultdict(list)\n",
    "        for x, y in edges:\n",
    "            e[x].append(y)\n",
    "            e[y].append(x)\n",
    "        s = set([(x, y) for x, y in guesses])\n",
    "        \n",
    "        cnt = 0\n",
    "        def dfs(node, pn):\n",
    "            nonlocal cnt\n",
    "            for ne in e[node]:\n",
    "                if ne == pn:\n",
    "                    continue\n",
    "                cnt += ((node, ne) in s)\n",
    "                dfs(ne, node)\n",
    "        dfs(0, -1)\n",
    "        q = deque()\n",
    "        q.append((0, cnt))\n",
    "        seen = set()\n",
    "        # seen.add(0)\n",
    "        res = 0\n",
    "        # print(cnt)\n",
    "        if cnt >= k:\n",
    "            res += 1\n",
    "        while q:\n",
    "            tmp, c = q.popleft()\n",
    "            seen.add(tmp)\n",
    "            # print(tmp, c)\n",
    "            for ne in e[tmp]:\n",
    "                if ne not in seen:\n",
    "                    new_c = c + ((ne, tmp) in s) - ((tmp, ne) in s)\n",
    "                    # print((ne, tmp) in s)\n",
    "                    # print((tmp, ne) in s)\n",
    "                    if new_c >= k:\n",
    "                        res += 1\n",
    "                    q.append((ne, new_c))\n",
    "        return res\n",
    "            \n",
    "                \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = max([max(e) for e in edges])+1\n",
    "        g = [[] for _ in range(n)]\n",
    "\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        gs = set([])\n",
    "        for f,c in guesses:\n",
    "            gs.add((f,c))\n",
    "        \n",
    "        def dfs(u,p):\n",
    "            res = 0\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    if (u,v) in gs:\n",
    "                        res += 1\n",
    "                    res += dfs(v,u)\n",
    "            return res\n",
    "        \n",
    "        tt = dfs(0,-1)\n",
    "        res = 0\n",
    "        def reroot(u,p,K):\n",
    "            nonlocal res\n",
    "            if K >= k:\n",
    "                res += 1\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    if ((u,v) in gs and (v,u) in gs) or ((u,v) not in gs and (v,u) not in gs):\n",
    "                        reroot(v,u,K)\n",
    "                    elif (u,v) in gs:\n",
    "                        reroot(v,u,K-1)\n",
    "                    else:\n",
    "                        reroot(v,u,K+1)\n",
    "        reroot(0,-1,tt)\n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = max([max(e) for e in edges])+1\n",
    "        g = [[] for _ in range(n)]\n",
    "\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        gs = set([])\n",
    "        for f,c in guesses:\n",
    "            gs.add((f,c))\n",
    "        \n",
    "        #返回以u为根的树中有几个正确的guess\n",
    "        def dfs(u,p):\n",
    "            res = 0\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    if (u,v) in gs:\n",
    "                        res += 1\n",
    "                    res += dfs(v,u)\n",
    "            return res\n",
    "        \n",
    "        #初始有K个正确guess\n",
    "        def reroot(u,p,K):\n",
    "            res = 0\n",
    "            if K >= k:\n",
    "                res = 1\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    #换根后正确的guess不变\n",
    "                    if ((u,v) in gs and (v,u) in gs) or ((u,v) not in gs and (v,u) not in gs):\n",
    "                        res += reroot(v,u,K)\n",
    "                    #换根后正确的guess少一个\n",
    "                    elif (u,v) in gs:\n",
    "                        res += reroot(v,u,K-1)\n",
    "                    #换根后正确的guess多一个\n",
    "                    else:\n",
    "                        res += reroot(v,u,K+1)\n",
    "            return res\n",
    "        return reroot(0,-1,dfs(0,-1))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# class Solution:\n",
    "#     def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "#         n = len(edges) + 1\n",
    "#         adj = collections.defaultdict(list)\n",
    "#         for a, b in edges:\n",
    "#             adj[a].append(b)\n",
    "#             adj[b].append(a)\n",
    "#         # hasht = dict()\n",
    "#         s = {(x, y) for x, y in guesses}\n",
    "#         ans = [0] * n\n",
    "#         def dfs1(node, root):\n",
    "#             for target in adj[node]:\n",
    "#                 if target != root:\n",
    "#                     if [node, target] in guesses:\n",
    "#                         ans[0] += 1\n",
    "#                     #     hasht[str(node) + '_' + str(target)] = 1\n",
    "#                     # else:\n",
    "#                     #     hasht[str(node) + '_' + str(target)] = 0\n",
    "#                     # if [target, node] in guesses:\n",
    "#                     #     hasht[str(target) + '_' + str(node)] = 1\n",
    "#                     # else:\n",
    "#                     #     hasht[str(target) + '_' + str(node)] = 0\n",
    "#                     dfs1(target, node)\n",
    "#         dfs1(0, -1)\n",
    "#         # print(hasht)\n",
    "#         def reroot(node, root):\n",
    "#             for target in adj[node]:\n",
    "#                 if target != root:\n",
    "#                     ans[target] = ans[node]\n",
    "#                     # if hasht[str(node) + '_' + str(target)]:\n",
    "#                     #     ans[target] -= 1\n",
    "#                     # if hasht[str(target) + '_' + str(node)]:\n",
    "#                     #     ans[target] += 1\n",
    "#                     if (node, target) in s:\n",
    "#                         ans[target] -= 1\n",
    "#                     if (target, node) in s:\n",
    "#                         ans[target] += 1\n",
    "#                     reroot(target, node)\n",
    "#         reroot(0, -1)\n",
    "#         # print(ans)\n",
    "#         return len([x for x in ans if x >= k])\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        g = [[] for _ in range(len(edges) + 1)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # s = {(x, y) for x, y in guesses}  # guesses 转成哈希表 s\n",
    "        s = set(map(tuple, guesses))\n",
    "        print(s)\n",
    "        ans = cnt0 = 0\n",
    "        def dfs(x: int, fa: int) -> None:\n",
    "            nonlocal cnt0\n",
    "            for y in g[x]:\n",
    "                if y != fa:\n",
    "                    cnt0 += (x, y) in s  # 以 0 为根时，猜对了\n",
    "                    dfs(y, x)\n",
    "        dfs(0, -1)\n",
    "\n",
    "        def reroot(x: int, fa: int, cnt: int) -> None:\n",
    "            nonlocal ans\n",
    "            ans += cnt >= k  # 此时 cnt 就是以 x 为根时的猜对次数\n",
    "            for y in g[x]:\n",
    "                if y != fa:\n",
    "                    reroot(y, x, cnt - ((x, y) in s) + ((y, x) in s))\n",
    "        reroot(0, -1, cnt0)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        guessList = [0] * n\n",
    "        sets={(a,b) for a,b in guesses}\n",
    "\n",
    "        def dfs(x, fa):\n",
    "            for y in g[x]:\n",
    "                if y != fa:\n",
    "                    if (x,y) in sets:\n",
    "                        guessList[0] += 1\n",
    "                    dfs(y, x)\n",
    "\n",
    "        dfs(0, -1)\n",
    "        # print(guessList)\n",
    "        def reroot(x, fa):\n",
    "            for y in g[x]:\n",
    "                if y != fa:\n",
    "                    guessList[y] = guessList[x]\n",
    "                    # print(guessList[y])\n",
    "                    if (x,y) in sets:\n",
    "                        guessList[y] -= 1\n",
    "                    if (y,x) in sets:\n",
    "                        guessList[y] =guessList[y]+1\n",
    "                    reroot(y,x)\n",
    "\n",
    "        reroot(0, -1)\n",
    "        # print(guessList)\n",
    "\n",
    "        return len([x for x in guessList if x>=k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        guesses = set([(u, v) for u, v in guesses])\n",
    "\n",
    "        self.count = 0\n",
    "        def dfs(node, parent):\n",
    "            for nei in g[node]:\n",
    "                if nei != parent:\n",
    "                    if (node, nei) in guesses:\n",
    "                        self.count += 1\n",
    "                    dfs(nei, node)\n",
    "        dfs(0, -1)\n",
    "\n",
    "        self.res = 0\n",
    "        def change_root(node, parent, count):\n",
    "            if count >= k:\n",
    "                self.res += 1\n",
    "            for nei in g[node]:\n",
    "                if nei != parent:\n",
    "                    change_root(nei, node, count - ((node, nei) in guesses) + ((nei, node) in guesses))\n",
    "        change_root(0, -1, self.count)\n",
    "        return self.res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        #e=defaultdict(set)\n",
    "\n",
    "        e=[[]for _ in range(len(edges)+1)]\n",
    "        for i in edges:\n",
    "            e[i[0]].append(i[1])\n",
    "            e[i[1]].append(i[0])\n",
    "        print(e)\n",
    "        s= {(x, y) for x, y in guesses} \n",
    "\n",
    "        def dfs(i,f):\n",
    "            nonlocal num_k\n",
    "            if (f,i) in s:\n",
    "                num_k+=1\n",
    "            for j in e[i]:\n",
    "                if j!=f:\n",
    "                    dfs(j,i)\n",
    "        n=len(edges)+1\n",
    "        res=0\n",
    "        num_k=0\n",
    "        dfs(0,-1)\n",
    "        num=[0]*n\n",
    "        num[0]=num_k\n",
    "\n",
    "        def fnew(i,f):\n",
    "            nonlocal res\n",
    "            for j in e[i]:\n",
    "                if j!=f:\n",
    "                    #print(j)\n",
    "                    num[j]=num[i]\n",
    "                    if (i,j) in s:\n",
    "                        num[j]-=1\n",
    "                    if (j,i) in s:\n",
    "                        num[j]+=1\n",
    "                    if num[j]>=k:\n",
    "                        res+=1\n",
    "                    fnew(j,i)\n",
    "        if num[0]>=k:\n",
    "            res=1\n",
    "        else:\n",
    "            res=0\n",
    "        fnew(0,-1)\n",
    "        \n",
    "        #print(num)\n",
    "        return res\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        aux = set()    \n",
    "        for edge in guesses:\n",
    "            aux.add(tuple(edge))\n",
    "        def dfs(u,fa)    :\n",
    "            cur = 0\n",
    "            for v in g[u]:\n",
    "                if v!= fa:\n",
    "                    cur+=dfs(v,u)\n",
    "                    if (u,v) in aux:\n",
    "                        cur +=1\n",
    "            return cur\n",
    "        n = len(edges) + 1\n",
    "        ans =[0] * n\n",
    "        ans[0] = dfs(0,-1)\n",
    "        def dfs2(u,fa):\n",
    "            for v in g[u]:\n",
    "                if v!= fa:\n",
    "                    ans[v] = ans[u]\n",
    "                    if (u,v) in aux:\n",
    "                        ans[v] -=1\n",
    "                    if (v,u)  in aux:   \n",
    "                        ans[v] += 1\n",
    "                    dfs2(v,u)\n",
    "        dfs2(0,-1)\n",
    "        return sum([1 for v in ans if v >= k])\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "\n",
    "        guesses = set([(u, v) for u, v in guesses])\n",
    "        cnt = Counter()\n",
    "\n",
    "        def dfs(node, pre):\n",
    "            cnt[node] += (pre, node) in guesses\n",
    "            for ch in g[node]:\n",
    "                if ch != pre:\n",
    "                    dfs(ch, node)\n",
    "                    cnt[node] += cnt[ch]\n",
    "\n",
    "        dfs(0, -1)\n",
    "        \n",
    "        res = [0] * n\n",
    "        res[0] = cnt[0]\n",
    "        def reroot(node, pre):\n",
    "            for ch in g[node]:\n",
    "                if ch != pre:\n",
    "                    res[ch] = res[node] + ((ch, node) in guesses) - ((node, ch) in guesses)\n",
    "                    reroot(ch, node)\n",
    "\n",
    "        reroot(0, -1)\n",
    "        return sum(x >= k for x in res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        st = set((u, v) for u, v in guesses)\n",
    "        g = defaultdict(list)\n",
    "        n = len(edges) + 1\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "            \n",
    "        def dfs(pa, x):\n",
    "            for y in g[x]:\n",
    "                if y != pa:\n",
    "                    if (x, y) in st:\n",
    "                        cnt[0] += 1\n",
    "                    dfs(x, y)\n",
    "        \n",
    "        def reroot(pa, x):\n",
    "            for y in g[x]:\n",
    "                if y != pa:\n",
    "                    if cnt[y] == 0:\n",
    "                        cnt[y] = cnt[x]\n",
    "                    if (x, y) in st:\n",
    "                        cnt[y] -= 1\n",
    "                    if (y, x) in st:\n",
    "                        cnt[y] += 1\n",
    "                    reroot(x, y)\n",
    "            \n",
    "        cnt = Counter()\n",
    "        dfs(-1, 0)\n",
    "        # print(f'{cnt}')\n",
    "\n",
    "        reroot(-1, 0)\n",
    "        # print(f'{cnt}')\n",
    "\n",
    "        return sum(1 for _, v in cnt.items() if v >= k)\n",
    "        \n",
    "            \n",
    "            \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for i, j in edges:\n",
    "            g[i].append(j)\n",
    "            g[j].append(i)\n",
    "        cnt = defaultdict(int)\n",
    "        for i, j in guesses:\n",
    "            cnt[(i, j)] += 1\n",
    "        \n",
    "        correct = 0\n",
    "        def dfs1(u, fa):\n",
    "            nonlocal correct\n",
    "            for nxt in g[u]:\n",
    "                if nxt == fa:\n",
    "                    continue\n",
    "                correct += cnt[(u, nxt)]\n",
    "                dfs1(nxt, u)\n",
    "        \n",
    "        ans = 0\n",
    "        def dfs2(u, fa):\n",
    "            nonlocal ans\n",
    "            nonlocal correct\n",
    "            ans += (correct >= k)\n",
    "            for nxt in g[u]:\n",
    "                if nxt == fa:\n",
    "                    continue\n",
    "                correct -= cnt[(u, nxt)]\n",
    "                correct += cnt[(nxt, u)]\n",
    "                dfs2(nxt, u)\n",
    "                correct -= cnt[(nxt, u)]\n",
    "                correct += cnt[(u, nxt)]\n",
    "\n",
    "        dfs1(0, -1)\n",
    "        dfs2(0, -1)\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        # 暴力解法：枚举每个根，得到father列表\n",
    "        # 方法2: 求得一个二维数组dis[i][j]代表i到j的最近距离 --> dijkstra: O(n*mlogm), 然后判断： O(n*k)\n",
    "        # 方法3:\n",
    "        # 注意到，如果节点 x 和 y 之间有边，那么从「以 x 为根的树」变成「以 y 为根的树」，就只有 [x,y] 和 [y,x] 这两个猜测的正确性变了，其余猜测的正确性不变。\n",
    "        # 先以0为根遍历一遍，得到father列表\n",
    "        # 再遍历一遍，每次遍历的时候换根\n",
    "        graph = collections.defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            graph[a].append(b)\n",
    "            graph[b].append(a)\n",
    "            \n",
    "        guss_fa = collections.defaultdict(list)\n",
    "        for fa, son in guesses:\n",
    "            guss_fa[son].append(fa)\n",
    "        # print(guss_fa)\n",
    "        # father列表记录以0为根的真值\n",
    "        father = dict()\n",
    "        count = 0\n",
    "        def dfs(node, fa):\n",
    "            nonlocal count\n",
    "            father[node] = fa\n",
    "            count += guss_fa[node].count(fa)\n",
    "            for son in graph[node]:\n",
    "                if son != fa:\n",
    "                    dfs(son, node)\n",
    "        dfs(0, -1)\n",
    "        # print(count)\n",
    "        ans = 0\n",
    "        def helper(node, count):\n",
    "            nonlocal ans\n",
    "            if count >= k:\n",
    "                ans += 1\n",
    "            for son in graph[node]:\n",
    "                if son != father[node]:\n",
    "                    # son 和 node调换位置\n",
    "                    temp = count\n",
    "                    father[node] = son\n",
    "                    temp += guss_fa[node].count(son)\n",
    "                    temp -= guss_fa[son].count(node)\n",
    "                    helper(son, temp) \n",
    "                    father[son] = node\n",
    "        helper(0, count)\n",
    "        return ans \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        st = set([tuple(x) for x in guesses])\n",
    "        dfd = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            dfd[a].append(b)\n",
    "            dfd[b].append(a)\n",
    "        f = [[0] * 2 for _ in range(n)]\n",
    "        def dfs_d(u, fa):\n",
    "            for ne in dfd[u]:\n",
    "                if ne != fa:\n",
    "                    dfs_d(ne, u)\n",
    "                    f[u][0] += f[ne][0] + ((u, ne) in st)\n",
    "        def dfs_u(u, fa):\n",
    "            if fa != -1:\n",
    "                f[u][1] = f[fa][0] - f[u][0] - ((fa, u) in st) + f[fa][1] + ((u, fa) in st)\n",
    "            for ne in dfd[u]:\n",
    "                if ne != fa:\n",
    "                    dfs_u(ne, u)\n",
    "        dfs_d(0, -1)\n",
    "        dfs_u(0, -1)\n",
    "        res = 0\n",
    "        for i in range(n):\n",
    "            if f[i][0] + f[i][1] >= k: res += 1\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n=len(edges)+1\n",
    "        d=defaultdict(int)\n",
    "        g=[[] for _ in range(n)]\n",
    "        for i,j in edges:\n",
    "            g[i].append(j)\n",
    "            g[j].append(i)\n",
    "        for i,j in guesses:\n",
    "            d[(i,j)]+=1\n",
    "        def dfs(x,fa):\n",
    "            res=0\n",
    "            for i in g[x]:\n",
    "                if i!=fa:\n",
    "                    res+=d[(x,i)]\n",
    "                    res+=dfs(i,x)\n",
    "            return res\n",
    "        ans=[0]*n\n",
    "        ans[0]=dfs(0,0)\n",
    "        def reroot(x,fa):\n",
    "            for i in g[x]:\n",
    "                if i!=fa:\n",
    "                    ans[i]=ans[x]+d[(i,x)]-d[(x,i)]\n",
    "                    reroot(i,x)\n",
    "        reroot(0,0)\n",
    "        return sum(x>=k for x in ans)\n",
    "\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [{} for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u][v] = 0\n",
    "            g[v][u] = 0\n",
    "\n",
    "        for u, v in guesses:\n",
    "            g[u][v] = 1\n",
    "\n",
    "        dic = {}\n",
    "\n",
    "        def dfs(node, pre):\n",
    "            res = 0\n",
    "            for ch in g[node]:\n",
    "                if ch != pre:\n",
    "                    res += dfs(ch, node)\n",
    "\n",
    "            dic[node] = res\n",
    "            if pre != -1:\n",
    "                res += g[pre][node]\n",
    "            return res\n",
    "\n",
    "        dfs(0, -1)\n",
    "\n",
    "        def dfs1(node, pre, cnt):\n",
    "            ans = 0\n",
    "            for ch in g[node]:\n",
    "                if ch != pre:\n",
    "                    t = cnt + g[ch][node]\n",
    "                    for c in g[node]:\n",
    "                        if c != pre and c != ch:\n",
    "                            t += dic[c] + g[node][c]\n",
    "                    ans += dfs1(ch, node, t)\n",
    "\n",
    "            if dic[node] + cnt >= k:\n",
    "                ans += 1\n",
    "            return ans\n",
    "\n",
    "        return dfs1(0, -1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        n=len(edges)+1\n",
    "        vis=set()\n",
    "        for u,v in guesses:\n",
    "            vis.add((u,v))\n",
    "        \n",
    "        g=defaultdict(list)\n",
    "        res=[0]*n\n",
    "        for u,v in edges:\n",
    "            g[u].append((v,int((u,v) in vis)))\n",
    "            g[v].append((u,int((v,u) in vis)))\n",
    "        \n",
    "        def dfs(u,fa):\n",
    "\n",
    "            for v,w in g[u]:\n",
    "                if v==fa:continue\n",
    "                res[0]+=w\n",
    "                dfs(v,u)\n",
    "        \n",
    "\n",
    "        def dfs2(u,fa):\n",
    "\n",
    "            for v,w in g[u]:\n",
    "                if v==fa:continue\n",
    "                res[v]=res[u]-w + int((v,u) in vis)\n",
    "                dfs2(v,u)\n",
    "            \n",
    "        dfs(0,-1)\n",
    "        dfs2(0,-1)\n",
    "        print(res)\n",
    "        return sum(v>=k for v in res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for u , v in edges:\n",
    "            graph[u].append(v)\n",
    "            graph[v].append(u)\n",
    "\n",
    "        hashmap = defaultdict(lambda :defaultdict(bool))\n",
    "        for u , v in guesses:\n",
    "            hashmap[u][v] = True\n",
    "\n",
    "        ans = 0\n",
    "\n",
    "        def dfs(node,parent):\n",
    "            ans = 0\n",
    "            for nxt in graph[node]:\n",
    "                if nxt != parent:\n",
    "                    ans += dfs(nxt,node) + int(hashmap[node][nxt])\n",
    "            return ans\n",
    "\n",
    "        x = dfs(0,-1)\n",
    "        ans = 0\n",
    "\n",
    "        def dfs(node,parent,cnt):\n",
    "            nonlocal ans\n",
    "            if cnt >= k:\n",
    "                ans += 1\n",
    "            for nxt in graph[node]:\n",
    "                if nxt != parent:\n",
    "                    dfs(nxt,node,cnt + int(hashmap[nxt][node]) - int(hashmap[node][nxt]))\n",
    "\n",
    "        dfs(0,-1,x)\n",
    "\n",
    "        return ans"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
